// This is the code for the BNs inference
//
//  'The second project homework for Artifical Intelligence'
// 
// Yangang Wang, 2020/12/2, @SEU
//

#include <iostream>
#include <fstream>
#include <string>
#include <regex>
#include <algorithm>
#include <vector>

using namespace std;

double k = 100;

vector<vector<vector<double>>> pxy(8, vector<vector<double>>(11, vector<double>(11, k)));
vector<vector<double>> px(8, vector<double>(11, 0));
vector<double> py(11, 0);

int Mat_missing[100][100];
int Mat_gd[100][100];
int Mat_ans[100][100];

int true_value_scale[11] = {0};

//  基本的设计思路就是计算文档中的贝叶斯公式，并选出概率最大的值填入-1的位置
//  px,py,pxy的含义以及公式推导见文档
//  2021.01.09  08017421-刘炜

//  计算px的值，px[i][j]代表xi=j的概率
void CalculatePx()
{
    for (int i = 1; i < 99; i++) {
        for (int j = 1; j < 99; j++) {
            if (Mat_missing[j - 1][i - 1] != -1) px[0][Mat_missing[j - 1][i - 1]] ++;
            if (Mat_missing[j - 1][i] != -1) px[1][Mat_missing[j - 1][i]] ++;
            if (Mat_missing[j - 1][i + 1] != -1) px[2][Mat_missing[j - 1][i + 1]] ++;
            if (Mat_missing[j][i - 1] != -1) px[3][Mat_missing[j][i - 1]] ++;
            if (Mat_missing[j][i + 1] != -1) px[4][Mat_missing[j][i + 1]] ++;
            if (Mat_missing[j + 1][i + 1] != -1) px[5][Mat_missing[j + 1][i + 1]] ++;
            if (Mat_missing[j + 1][i] != -1) px[6][Mat_missing[j + 1][i]] ++;
            if (Mat_missing[j + 1][i - 1] != -1) px[7][Mat_missing[j + 1][i - 1]] ++;
        }
    }

    for (int i = 0; i < 8; i++) {
        int sum = 0;
        for (int j = 0; j < 11; j++)
            sum += px[i][j];
        for (int j = 0; j < 11; j++)
            px[i][j] /= sum;
    }
}

//  计算py的值，py[i]代表y=i的概率
void CalculatePy()
{
    for (int i = 1; i < 99; i++) {
        for (int j = 1; j < 99; j++) {
            if (Mat_missing[j][i] == -1) continue;
            py[Mat_missing[j][i]] ++;
        }
    }

    int sum = 0;
    for (int i = 0; i < 11; i++)
        sum += py[i];
    for (int i = 0; i < 11; i++)
        py[i] /= sum;
}

//  计算pxy的值，pxy[i][j][k]代表已知y=k的情况下，xi=j的概率
void CalculatePxy()
{
    for (int i = 1; i < 99; i++) {
        for (int j = 1; j < 99; j++) {
            if (Mat_missing[j][i] == -1) continue;
            if (Mat_missing[j - 1][i - 1] != -1) pxy[0][Mat_missing[j - 1][i - 1]][Mat_missing[j][i]] ++;
            if (Mat_missing[j - 1][i] != -1) pxy[1][Mat_missing[j - 1][i]][Mat_missing[j][i]] ++;
            if (Mat_missing[j - 1][i + 1] != -1) pxy[2][Mat_missing[j - 1][i + 1]][Mat_missing[j][i]] ++;

            if (Mat_missing[j][i - 1] != -1) pxy[3][Mat_missing[j][i - 1]][Mat_missing[j][i]] ++;
            if (Mat_missing[j][i + 1] != -1) pxy[4][Mat_missing[j][i + 1]][Mat_missing[j][i]] ++;

            if (Mat_missing[j + 1][i + 1] != -1) pxy[5][Mat_missing[j + 1][i + 1]][Mat_missing[j][i]] ++;
            if (Mat_missing[j + 1][i] != -1) pxy[6][Mat_missing[j + 1][i]][Mat_missing[j][i]] ++;
            if (Mat_missing[j + 1][i - 1] != -1) pxy[7][Mat_missing[j + 1][i - 1]][Mat_missing[j][i]] ++;
        }
    }

    for (int i = 0; i < 8; i++) {
        for (int j = 0; j < 11; j++) {
            int sum = 11 * k;
            for (int m = 0; m < 11; m++)
                sum += pxy[i][m][j];
            for (int m = 0; m < 11; m++)
                pxy[i][m][j] /= sum;
        }
    }
}

// 预测函数
void Predict()
{
    int bigger[102][102] = { -1 };
    for (int i = 0; i < 100; i++) {
        for (int j = 0; j < 100; j++) {
            bigger[i + 1][j + 1] = Mat_missing[i][j];
            Mat_ans[i][j] = Mat_missing[i][j];
        }
    }
    CalculatePx();
    CalculatePy();
    CalculatePxy();

    vector<int> x(8);
    for (int i = 1; i < 101; i++) {
        for (int j = 1; j < 101; j++) {
            if (bigger[i][j] != -1) continue;
            vector<double> p(11, 1.0);
            {
                x[0] = bigger[i - 1][j - 1];
                x[1] = bigger[i - 1][j];
                x[2] = bigger[i - 1][j + 1];
                x[3] = bigger[i][j - 1];
                x[4] = bigger[i][j + 1];
                x[5] = bigger[i + 1][j - 1];
                x[6] = bigger[i + 1][j];
                x[7] = bigger[i + 1][j + 1];
            }
            
            double x_avg = 0;
            for (int i = 0; i < 8; i++)
                x_avg += x[i] < 0 ? 0 : x[i];
            x_avg /= 8;
            for (int i = 0; i < 8; i++)
                x[i] = x[i] < 0 ? x_avg : x[i];

            for (int n = 0; n < 11; n++) {
                p[n] = py[n];
                for (int i = 0; i < 8; i++)
                    p[n] *= pxy[i][x[i]][n] / px[i][x[i]];
            }
            auto biggest = std::max_element(std::begin(p), std::end(p));
            auto ans = biggest - p.begin();
            Mat_ans[i - 1][j - 1] = (int)ans;
        }
    }
}

// 计算正确率
void CalculateAccuarcy()
{
    int total = 0;
    int wrong = 0;
    for (int i = 0; i < 100; i++) {
        for (int j = 0; j < 100; j++) {
            if (Mat_missing[i][j] == -1) {
                total++;
                true_value_scale[Mat_gd[i][j]]++;
                if (Mat_ans[i][j] != Mat_gd[i][j])
                    wrong++;
            }
        }
    }

    cout << "wrong: " << wrong << endl;
    cout << "total: " << total << endl;
    double accuarcy = 0.0;
    accuarcy = 1.0 * (total - wrong) / total;
    cout << "accuarcy: " << accuarcy << endl;
    for (int i = 0; i <= 10; i++) {
        cout << "num of " << i << " : " << true_value_scale[i] << endl;
    }
}

int main(int argc, char** argv)
{
	// read the mat
	// the missing element is indicated by -1
	//
	{
		ifstream infile("mat_missing.txt");
		if (infile.is_open()) {
			string line;
			regex term(" ");
			int lineIdx = 0;
			while (!infile.eof()) {
				if (lineIdx >= 100) break;
				getline(infile, line);
				sregex_token_iterator it(line.begin(), line.end(), term, -1);
				sregex_token_iterator end;
				for (int i = 0; it != end; it++, i++) {
					string str = *it;
					Mat_missing[lineIdx][i] = atoi(str.c_str());
				}
				lineIdx++;
			}
			infile.close();
		}
	}

	//======================
	// please fill the following code by Bayes Network training and inference
	//
	// Yangang Wang, 2020/12/2, @SEU
	//

	// (1) todo...
    Predict();
	//======================
	// You can compare the result with groudtruth data
	// read the groudtruth data
	{
		ifstream infile("mat_groudtruth.txt");
		if (infile.is_open()) {
			string line;
			regex term(" ");
			int lineIdx = 0;
			while (!infile.eof()) {
				if (lineIdx >= 100) break;
				getline(infile, line);
				sregex_token_iterator it(line.begin(), line.end(), term, -1);
				sregex_token_iterator end;
				for (int i = 0; it != end; it++, i++) {
					string str = *it;
					Mat_gd[lineIdx][i] = atoi(str.c_str());
				}
				lineIdx++;
			}
			infile.close();
		}
	}
	// fill the comparing code here
	//
	// Yangang Wang, 2020/12/2, @SEU
	//
    CalculateAccuarcy();
	// (2) todo...
    return 0;
}
